{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "assignment_9_Zheyu_Xie_5layers_trainFMNIST.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "0zklGxKVx5Y_",
        "colab_type": "code",
        "outputId": "c8ccc65d-760a-4ea5-a212-087b85347f07",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 125
        }
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bKB1z7JJPftR",
        "colab_type": "text"
      },
      "source": [
        "This is the only cell to be modified"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yy09ePlLejit",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import sys\n",
        "import re\n",
        "import time\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "from datetime import datetime\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.init as init\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data.sampler import *\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision import transforms, datasets\n",
        "\n",
        "class A9_Params:\n",
        "    \"\"\"\n",
        "    :ivar dataset:\n",
        "        0: MNIST\n",
        "        1: FMNIST\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        self.use_cuda = 1\n",
        "        self.dataset = 1\n",
        "        self.train_split = 0.8\n",
        "        self.labeled_split = 0.2\n",
        "        self.mnist = MNISTParams()\n",
        "        self.fmnist = FMNISTParams()\n",
        "\n",
        "class TrainParams:\n",
        "    \"\"\"\n",
        "    :ivar optim_type:\n",
        "      0: SGD\n",
        "      1: ADAM\n",
        "\n",
        "    :ivar load_weights:\n",
        "        0: train from scratch\n",
        "        1: load and test\n",
        "        2: load if it exists and continue training\n",
        "\n",
        "    :ivar save_criterion:  when to save a new checkpoint\n",
        "        0: max validation accuracy\n",
        "        1: min validation loss\n",
        "        2: max training accuracy\n",
        "        3: min training loss\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        self.batch_size = 128\n",
        "        self.optim_type = 1\n",
        "        \n",
        "        #self.lr = 0.00005 #before 100 epoc\n",
        "        self.lr = 0.00003\n",
        "\n",
        "        self.momentum = 0.9\n",
        "        self.n_epochs = 400\n",
        "        self.weight_decay = 0.00008\n",
        "        self.c0 = 0.35\n",
        "        self.save_criterion = 0\n",
        "        self.load_weights = 1\n",
        "        self.weights_path = '/content/drive/My Drive/Colab Notebooks/a9/model.pt'\n",
        "\n",
        "\n",
        "class MNISTParams(TrainParams):\n",
        "    def __init__(self):\n",
        "        super(MNISTParams, self).__init__()\n",
        "        #self.weights_path = '/content/drive/My Drive/A9/checkpoints/mnist/model.pt'\n",
        "        \n",
        "        self.weights_path = '/content/drive/My Drive/Colab Notebooks/checkpoints/mnist'\n",
        "\n",
        "\n",
        "class FMNISTParams(TrainParams):\n",
        "    def __init__(self):\n",
        "        super(FMNISTParams, self).__init__()\n",
        "        #self.weights_path = '/content/drive/My Drive/A9/checkpoints/fmnist/model.pt'\n",
        "        self.weights_path = '/content/drive/My Drive/Colab Notebooks/a9/checkpoints/fmnist'\n",
        "\n",
        "class MNISTParams(TrainParams):\n",
        "    def __init__(self):\n",
        "        super(MNISTParams, self).__init__()\n",
        "        #self.weights_path = './checkpoints/mnist/model.pt'\n",
        "        \n",
        "        self.weights_path = '/content/drive/My Drive/Colab Notebooks/checkpoints/mnist/model.pt'\n",
        "        \n",
        "\n",
        "class FMNISTParams(TrainParams):\n",
        "    def __init__(self):\n",
        "        super(FMNISTParams, self).__init__()\n",
        "        #self.weights_path = './checkpoints/fmnist/model.pt'\n",
        "        self.weights_path = '/content/drive/My Drive/Colab Notebooks/checkpoints/fmnist/model.pt'\n",
        "\n",
        "\n",
        "class CompositeLoss(nn.Module):\n",
        "    def __init__(self, device):\n",
        "        super(CompositeLoss, self).__init__()\n",
        "        #self.reconstruction_loss = loss_rec\n",
        "        #self.classification_loss = loss_cls\n",
        "        pass\n",
        "\n",
        "    def init_weights(self):\n",
        "        pass\n",
        "\n",
        "    def forward(self, reconstruction_loss, classification_loss):\n",
        "        #loss = reconstruction_loss + c0 * classification_loss\n",
        "        #return loss\n",
        "        pass\n",
        "\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "  \n",
        "  #class variables\n",
        "  #weight1 = nn.Parameter(torch.rand(28*28, 512))\n",
        "  #weight2 = nn.Parameter(torch.rand(512,256))\n",
        "  #weight3 = nn.Parameter(torch.rand(256,256))\n",
        "  #weight4 = nn.Parameter(torch.rand(256,128))\n",
        "\n",
        "    def __init__(self, device):\n",
        "      super(Encoder, self).__init__()\n",
        "      self.encod1 = torch.nn.Linear(28*28, 512, True)   #initize to extract weight \n",
        "      self.encod2 = torch.nn.Linear(512, 384, True)\n",
        "      self.encod3 = torch.nn.Linear(384, 384, True)\n",
        "      self.encod4 = torch.nn.Linear(384, 256, True)\n",
        "      self.encod5 = torch.nn.Linear(256, 150, True)\n",
        "\n",
        "      self.weight1 = nn.Parameter(self.encod1.weight)\n",
        "      self.weight2 = nn.Parameter(self.encod2.weight)\n",
        "      self.weight3 = nn.Parameter(self.encod3.weight)\n",
        "      self.weight4 = nn.Parameter(self.encod4.weight)\n",
        "      self.weight5 = nn.Parameter(self.encod5.weight)\n",
        "      \n",
        "      self.b1 = nn.Parameter(torch.randn(512)) #bias dim equal this layer's output dim\n",
        "      self.b2 = nn.Parameter(torch.randn(384))\n",
        "      self.b3 = nn.Parameter(torch.randn(384))\n",
        "      self.b4 = nn.Parameter(torch.randn(256))\n",
        "      self.b5 = nn.Parameter(torch.randn(150))\n",
        "\n",
        "    def get_weights(self):\n",
        "\n",
        "      #share weights between encoder and decoder, weight in decoder just transpose of corrsponding weight\n",
        "      return self.weight5.t(), self.weight4.t(), self.weight3.t(), self.weight2.t(), self.weight1.t()\n",
        "\n",
        "    def init_weights(self):\n",
        "      torch.nn.init.xavier_uniform_(self.encod1.weight)\n",
        "      torch.nn.init.xavier_uniform_(self.encod2.weight)\n",
        "      torch.nn.init.xavier_uniform_(self.encod3.weight)\n",
        "      torch.nn.init.xavier_uniform_(self.encod4.weight)\n",
        "      torch.nn.init.xavier_uniform_(self.encod5.weight)\n",
        "        \n",
        "\n",
        "    def forward(self, enc_input):\n",
        "      \n",
        "      x = enc_input\n",
        "      x = x.view(x.size(0), -1)\n",
        "    \n",
        "      x = F.elu(torch.mm(x, self.weight1.t()) + self.b1)\n",
        "      x = F.elu(torch.mm(x, self.weight2.t()) + self.b2)\n",
        "      x = F.elu(torch.mm(x, self.weight3.t()) + self.b3)\n",
        "      x = F.elu(torch.mm(x, self.weight4.t()) + self.b4)\n",
        "      x = F.elu(torch.mm(x, self.weight5.t()) + self.b5)\n",
        "      return x\n",
        "\n",
        "class Decoder(nn.Module):\n",
        "  \n",
        "  ''' \n",
        "self.dec1 is tensor inti by nn.Linear() so that can pass to Xe initiler\n",
        "self.dec_weight1 is a register paremeter, use for manully calculta y = xW' + b'''\n",
        "  \n",
        "  def __init__(self, device):\n",
        "    \n",
        "    super(Decoder, self).__init__()\n",
        "    self.dec_weight1 = torch.randn(150,256)\n",
        "    self.dec_weight2 = torch.randn(256,384)\n",
        "    self.dec_weight3 = torch.randn(384,384)\n",
        "    self.dec_weight4 = torch.randn(384,512)\n",
        "    self.dec_weight5 = torch.randn(512,28*28)\n",
        "    \n",
        "    self.dec_b1 = nn.Parameter(torch.randn(256))  #dim= output dim of that layer\n",
        "    self.dec_b2 = nn.Parameter(torch.randn(384))\n",
        "    self.dec_b3 = nn.Parameter(torch.randn(384))\n",
        "    self.dec_b4 = nn.Parameter(torch.randn(512))\n",
        "    self.dec_b5 = nn.Parameter(torch.randn(28*28))\n",
        "\n",
        "  def init_weights(self, shared_weights):\n",
        "    \n",
        "    #print(type(self.dec_weight1))\n",
        "    #print(type(shared_weights[0]))\n",
        "    self.dec_weight1, self.dec_weight2, self.dec_weight3, self.dec_weight4, self.dec_weight5 = shared_weights\n",
        "\n",
        "    #convert to parameters\n",
        "    self.dec_weight1 = nn.Parameter(self.dec_weight1)\n",
        "    self.dec_weight2 = nn.Parameter(self.dec_weight2)\n",
        "    self.dec_weight3 = nn.Parameter(self.dec_weight3)\n",
        "    self.dec_weight4 = nn.Parameter(self.dec_weight4)\n",
        "    self.dec_weight5 = nn.Parameter(self.dec_weight5)\n",
        "\n",
        "  def forward(self, dec_input):\n",
        "    x = dec_input\n",
        "    #print(x.shape)\n",
        "    #print(self.dec_weight1.t().shape)\n",
        "    #print(self.dec_b1.shape)\n",
        "    x = F.elu(torch.mm(x, self.dec_weight1.t()) + self.dec_b1)\n",
        "    x = F.elu(torch.mm(x, self.dec_weight2.t()) + self.dec_b2)\n",
        "    x = F.elu(torch.mm(x, self.dec_weight3.t()) + self.dec_b3)\n",
        "    x = F.elu(torch.mm(x, self.dec_weight4.t()) + self.dec_b4)\n",
        "    x = F.elu(torch.mm(x, self.dec_weight5.t()) + self.dec_b5)\n",
        "\n",
        "    x = x.view(x.size(0), 1, 28, 28)\n",
        "    return x\n",
        "\n",
        "\n",
        "class Classifier(nn.Module):\n",
        "  \n",
        "  \n",
        "  def __init__(self, device):\n",
        "    \n",
        "    super(Classifier, self).__init__()\n",
        "    self.class_layer1 = torch.nn.Linear(150,150)\n",
        "    self.class_layer2 = torch.nn.Linear(150,150)\n",
        "    self.class_layer3 = torch.nn.Linear(150,100)\n",
        "    self.class_layer4 = torch.nn.Linear(100,60)\n",
        "    self.class_layer5 = torch.nn.Linear(60,10)\n",
        "\n",
        "\n",
        "  def init_weights(self):\n",
        "    \n",
        "    torch.nn.init.xavier_uniform_(self.class_layer1.weight)\n",
        "    torch.nn.init.xavier_uniform_(self.class_layer2.weight)\n",
        "    torch.nn.init.xavier_uniform_(self.class_layer3.weight)\n",
        "    torch.nn.init.xavier_uniform_(self.class_layer4.weight)\n",
        "    torch.nn.init.xavier_uniform_(self.class_layer5.weight)\n",
        "\n",
        "  def forward(self, x):\n",
        "    \n",
        "    x = x.view(x.size(0), -1)\n",
        "    x = F.relu(self.class_layer1(x))\n",
        "    x = F.relu(self.class_layer2(x))\n",
        "    x = F.relu(self.class_layer3(x))\n",
        "    x = F.relu(self.class_layer4(x))\n",
        "    x = self.class_layer5(x)\n",
        "    x = F.softmax(x,dim=1)\n",
        "    return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LDKEGAJgyKgh",
        "colab_type": "text"
      },
      "source": [
        "Helper code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5oXbKkpuPWvX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class PartiallyLabeled(Dataset):\n",
        "    \"\"\"\n",
        "    :param Dataset _dataset:\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dataset, all_idx, labeled_percent):\n",
        "        \"\"\"\n",
        "\n",
        "        :param Dataset dataset:\n",
        "        :param list all_idx:\n",
        "        :param float labeled_percent:\n",
        "        \"\"\"\n",
        "        self._dataset = dataset\n",
        "        self._n_data = len(all_idx)\n",
        "        self.n_labeled_data = int(labeled_percent * self._n_data)\n",
        "        self._is_labeled = np.zeros((self._n_data, 1), dtype=np.bool)\n",
        "        labeled_images = np.random.permutation(all_idx)[:self.n_labeled_data]\n",
        "        self._is_labeled[labeled_images] = 1\n",
        "\n",
        "    def __len__(self):\n",
        "        return self._dataset.__len__()\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        assert idx < self._n_data, \"Invalid idx: {} for _n_data: {}\".format(idx, self._n_data)\n",
        "\n",
        "        input, target = self._dataset.__getitem__(idx)\n",
        "        is_labeled = self._is_labeled[idx]\n",
        "        return input, target, is_labeled\n",
        "\n",
        "\n",
        "def get_psnr(x, x_test):\n",
        "    mse = np.mean((np.reshape(x_test, [-1, 28, 28]) - np.reshape(x, [-1, 28, 28])) ** 2)\n",
        "    psnr = -100.0 * np.log10(mse)\n",
        "    return psnr\n",
        "\n",
        "\n",
        "def eval(modules, data_loader, criteria, device):\n",
        "    modules.eval()\n",
        "    encoder, decoder, classifier = modules\n",
        "    criterion_rec, criterion_cls = criteria\n",
        "    mean_loss_sum = 0\n",
        "    _psnr_sum = 0\n",
        "    total = 0\n",
        "    correct = 0\n",
        "    n_batches = 0\n",
        "    with torch.no_grad():\n",
        "        for batch_idx, (inputs, targets) in enumerate(data_loader):\n",
        "            inputs = inputs.to(device)\n",
        "            targets = targets.to(device)\n",
        "\n",
        "            outputs_enc = encoder(inputs)\n",
        "            outputs_rec = decoder(outputs_enc)\n",
        "            outputs_cls = classifier(outputs_enc)\n",
        "\n",
        "            loss_rec = criterion_rec(outputs_rec, inputs)\n",
        "            loss_cls = criterion_cls(outputs_cls, targets)\n",
        "\n",
        "            loss = loss_rec + loss_cls\n",
        "\n",
        "            mean_loss = loss.item()\n",
        "\n",
        "            mean_loss_sum += mean_loss\n",
        "\n",
        "            _, predicted = outputs_cls.max(1)\n",
        "            total += targets.size(0)\n",
        "            correct += predicted.eq(targets).sum().item()\n",
        "\n",
        "            n_batches += 1\n",
        "\n",
        "            inputs_np = inputs.detach().cpu().numpy()\n",
        "            outputs_np = outputs_rec.detach().cpu().numpy()\n",
        "\n",
        "            _psnr = get_psnr(inputs_np, outputs_np)\n",
        "            _psnr_sum += _psnr\n",
        "\n",
        "    overall_mean_loss = mean_loss_sum / n_batches\n",
        "    mean_psnr = _psnr_sum / n_batches\n",
        "    acc = 100. * correct / total\n",
        "\n",
        "    return overall_mean_loss, acc, mean_psnr"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7AjqRf6TyEnb",
        "colab_type": "text"
      },
      "source": [
        "Main training code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JNIREOVbtbx1",
        "colab_type": "code",
        "outputId": "15595bec-779d-4802-c899-dfc876fd2887",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 512
        }
      },
      "source": [
        "params = A9_Params()\n",
        "\n",
        "# init device\n",
        "if params.use_cuda and torch.cuda.is_available():\n",
        "    device = torch.device(\"cuda\")\n",
        "    print('Training on GPU: {}'.format(torch.cuda.get_device_name(0)))\n",
        "else:\n",
        "    device = torch.device(\"cpu\")\n",
        "    print('Training on CPU')\n",
        "\n",
        "# load dataset\n",
        "if params.dataset == 0:\n",
        "    print('Using MNIST dataset')\n",
        "    transform = transforms.Compose([transforms.ToTensor(),\n",
        "                                    transforms.Normalize((0.1307,), (0.3081,)),\n",
        "                                    ])\n",
        "    train_set = datasets.MNIST('data', train=True, download=True, transform=transform)\n",
        "    test_set = datasets.MNIST('data', train=False, download=True, transform=transform)\n",
        "    valid_set = datasets.MNIST('data', train=True, download=True, transform=transform)\n",
        "    train_params = params.mnist\n",
        "elif params.dataset == 1:\n",
        "    print('Using Fashion MNIST dataset')\n",
        "    transform = transforms.Compose([transforms.ToTensor(),\n",
        "                                    transforms.Normalize((0.5,), (0.5,)),\n",
        "                                    ])\n",
        "    train_set = datasets.FashionMNIST('data', train=True, download=True, transform=transform)\n",
        "    test_set = datasets.FashionMNIST('data', train=False, download=True, transform=transform)\n",
        "    valid_set = datasets.FashionMNIST('data', train=True, download=True, transform=transform)\n",
        "    train_params = params.fmnist\n",
        "else:\n",
        "    raise IOError('Invalid db_type: {}'.format(params.dataset))\n",
        "\n",
        "num_train = len(train_set)\n",
        "indices = list(range(num_train))\n",
        "split = int(np.floor(params.train_split * num_train))\n",
        "\n",
        "train_idx, valid_idx = indices[:split], indices[split:]\n",
        "train_set = PartiallyLabeled(train_set, train_idx, labeled_percent=params.labeled_split)\n",
        "\n",
        "print('Training samples: {}\\n'\n",
        "      'Validation samples: {}\\n'\n",
        "      'Labeled training samples: {}'\n",
        "      ''.format(\n",
        "    len(train_idx),\n",
        "    len(valid_idx),\n",
        "    train_set.n_labeled_data\n",
        "))\n",
        "\n",
        "train_sampler = SubsetRandomSampler(train_idx)\n",
        "valid_sampler = SequentialSampler(valid_idx)\n",
        "\n",
        "train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=train_params.batch_size, sampler=train_sampler,\n",
        "                                                num_workers=4)\n",
        "valid_dataloader = torch.utils.data.DataLoader(valid_set, batch_size=24, sampler=valid_sampler,\n",
        "                                                num_workers=4)\n",
        "test_dataloader = torch.utils.data.DataLoader(test_set, batch_size=10, shuffle=False, num_workers=4)\n",
        "\n",
        "# create modules\n",
        "encoder = Encoder(device).to(device)\n",
        "decoder = Decoder(device).to(device)\n",
        "classifier = Classifier(device).to(device)\n",
        "\n",
        "assert isinstance(encoder, nn.Module), 'encoder must be an instance of nn.Module'\n",
        "assert isinstance(decoder, nn.Module), 'decoder must be an instance of nn.Module'\n",
        "assert isinstance(classifier, nn.Module), 'classifier must be an instance of nn.Module'\n",
        "\n",
        "modules = nn.ModuleList((encoder, decoder, classifier))\n",
        "\n",
        "# init weights\n",
        "encoder.init_weights()\n",
        "decoder.init_weights(encoder.get_weights())\n",
        "classifier.init_weights()\n",
        "\n",
        "# create losses\n",
        "criterion_rec = torch.nn.MSELoss().to(device)\n",
        "criterion_cls = torch.nn.CrossEntropyLoss().to(device)\n",
        "\n",
        "parameters = list(modules.parameters())\n",
        "if train_params.c0 == 0:\n",
        "    composite_loss = CompositeLoss(device)\n",
        "    composite_loss.init_weights()\n",
        "    assert isinstance(composite_loss, nn.Module), 'composite_loss must be an instance of nn.Module'\n",
        "    parameters += list(composite_loss.parameters())\n",
        "else:\n",
        "    def composite_loss(x, y):\n",
        "        return x + train_params.c0 * y\n",
        "\n",
        "# create optimizer\n",
        "if train_params.optim_type == 0:\n",
        "    optimizer = torch.optim.SGD(parameters, lr=train_params.lr, momentum=train_params.momentum,\n",
        "                                weight_decay=train_params.weight_decay)\n",
        "elif train_params.optim_type == 1:\n",
        "    optimizer = torch.optim.Adam(parameters, lr=train_params.lr, weight_decay=train_params.weight_decay)\n",
        "else:\n",
        "    raise IOError('Invalid optim_type: {}'.format(train_params.optim_type))\n",
        "\n",
        "weights_dir = os.path.dirname(train_params.weights_path)\n",
        "weights_name = os.path.basename(train_params.weights_path)\n",
        "\n",
        "if not os.path.isdir(weights_dir):\n",
        "    os.makedirs(weights_dir)\n",
        "\n",
        "start_epoch = 0\n",
        "max_valid_acc_epoch = 0\n",
        "max_valid_acc = 0\n",
        "max_train_acc = 0\n",
        "min_valid_loss = np.inf\n",
        "min_train_loss = np.inf\n",
        "\n",
        "# load weights\n",
        "if train_params.load_weights:\n",
        "    matching_ckpts = [k for k in os.listdir(weights_dir) if\n",
        "                      os.path.isfile(os.path.join(weights_dir, k)) and\n",
        "                      k.startswith(weights_name)]\n",
        "    if not matching_ckpts:\n",
        "        msg = 'No checkpoints found matching {} in {}'.format(weights_name, weights_dir)\n",
        "        if train_params.load_weights == 1:\n",
        "            raise IOError(msg)\n",
        "        print(msg)\n",
        "    else:\n",
        "        matching_ckpts.sort(key=lambda x: [int(c) if c.isdigit() else c for c in re.split(r'(\\d+)', x)])\n",
        "\n",
        "        weights_path = os.path.join(weights_dir, matching_ckpts[-1])\n",
        "\n",
        "        chkpt = torch.load(weights_path, map_location=device)  # load checkpoint\n",
        "\n",
        "        print('Loading weights from: {} with:\\n'\n",
        "              '\\tepoch: {}\\n'\n",
        "              '\\ttrain_loss: {}\\n'\n",
        "              '\\ttrain_acc: {}\\n'\n",
        "              '\\tvalid_loss: {}\\n'\n",
        "              '\\tvalid_acc: {}\\n'\n",
        "              '\\ttimestamp: {}\\n'.format(\n",
        "            weights_path, chkpt['epoch'],\n",
        "            chkpt['train_loss'], chkpt['train_acc'],\n",
        "            chkpt['valid_loss'], chkpt['valid_acc'],\n",
        "            chkpt['timestamp']))\n",
        "\n",
        "        encoder.load_state_dict(chkpt['encoder'])\n",
        "        decoder.load_state_dict(chkpt['decoder'])\n",
        "        classifier.load_state_dict(chkpt['classifier'])\n",
        "        optimizer.load_state_dict(chkpt['optimizer'])\n",
        "\n",
        "        if train_params.c0 == 0 and 'composite_loss' in chkpt:\n",
        "            composite_loss.load_state_dict(chkpt['composite_loss'])\n",
        "\n",
        "        max_valid_acc = chkpt['valid_acc']\n",
        "        min_valid_loss = chkpt['valid_loss']\n",
        "\n",
        "        max_train_acc = chkpt['train_acc']\n",
        "        min_train_loss = chkpt['train_loss']\n",
        "\n",
        "        max_valid_acc_epoch = chkpt['epoch']\n",
        "        start_epoch = chkpt['epoch'] + 1\n",
        "\n",
        "if train_params.load_weights != 1:\n",
        "    # continue training\n",
        "    for epoch in range(start_epoch, train_params.n_epochs):\n",
        "        # Training\n",
        "        modules.train()\n",
        "\n",
        "        train_loss_rec = 0\n",
        "        train_loss_cls = 0\n",
        "        train_loss = 0\n",
        "        train_total = 0\n",
        "        train_correct = 0\n",
        "        batch_idx = 0\n",
        "\n",
        "        save_weights = 0\n",
        "\n",
        "        for batch_idx, (inputs, targets, is_labeled) in tqdm(enumerate(train_dataloader)):\n",
        "            inputs = inputs.to(device)\n",
        "            targets = targets.to(device)\n",
        "\n",
        "            if not np.count_nonzero(is_labeled.detach().numpy()):\n",
        "                continue\n",
        "                \n",
        "            is_labeled = is_labeled.squeeze().to(device)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            outputs_enc = encoder(inputs)\n",
        "            outputs_rec = decoder(outputs_enc)\n",
        "            outputs_cls = classifier(outputs_enc)\n",
        "\n",
        "            loss_rec = criterion_rec(outputs_rec, inputs)\n",
        "            loss_cls = criterion_cls(outputs_cls[is_labeled, :], targets[is_labeled])\n",
        "\n",
        "            loss = composite_loss(loss_rec, loss_cls)\n",
        "\n",
        "            mean_loss_rec = loss_rec.item()\n",
        "            mean_loss_cls = loss_cls.item()\n",
        "            train_loss_rec += mean_loss_rec\n",
        "            train_loss_cls += mean_loss_cls\n",
        "\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            mean_loss = loss.item()\n",
        "            train_loss += mean_loss\n",
        "\n",
        "            _, predicted = outputs_cls.max(1)\n",
        "            train_total += targets.size(0)\n",
        "            train_correct += predicted.eq(targets).sum().item()\n",
        "\n",
        "        mean_train_loss_rec = train_loss_rec / (batch_idx + 1)\n",
        "        mean_train_loss_cls = train_loss_cls / (batch_idx + 1)\n",
        "        mean_train_loss = train_loss / (batch_idx + 1)\n",
        "\n",
        "        train_acc = 100. * train_correct / train_total\n",
        "\n",
        "        valid_loss, valid_acc, valid_psnr = eval(\n",
        "            modules, valid_dataloader, (criterion_rec, criterion_cls), device)\n",
        "\n",
        "        if valid_acc > max_valid_acc:\n",
        "            max_valid_acc = valid_acc\n",
        "            max_valid_acc_epoch = epoch\n",
        "            if train_params.save_criterion == 0:\n",
        "                save_weights = 1\n",
        "\n",
        "        if valid_loss < min_valid_loss:\n",
        "            min_valid_loss = valid_loss\n",
        "            if train_params.save_criterion == 1:\n",
        "                save_weights = 1\n",
        "\n",
        "        if train_acc > max_train_acc:\n",
        "            max_train_acc = train_acc\n",
        "            if train_params.save_criterion == 2:\n",
        "                save_weights = 1\n",
        "\n",
        "        if train_loss < min_train_loss:\n",
        "            min_train_loss = train_loss\n",
        "            if train_params.save_criterion == 3:\n",
        "                save_weights = 1\n",
        "\n",
        "        print(\n",
        "            'Epoch: %d Train-Loss: %.6f (rec: %.6f, cls: %.6f) | Train-Acc: %.3f%% | '\n",
        "            'Validation-Loss: %.6f | Validation-Acc: %.3f%% | Validation-PSNR: %.3f | '\n",
        "            'Max Validation-Acc: %.3f%% (epoch: %d)' % (\n",
        "                epoch, mean_train_loss, mean_train_loss_rec, mean_train_loss_cls, train_acc,\n",
        "                valid_loss, valid_acc, valid_psnr, max_valid_acc, max_valid_acc_epoch))\n",
        "\n",
        "        # Save checkpoint.\n",
        "        if save_weights:\n",
        "            model_dict = {\n",
        "                'encoder': encoder.state_dict(),\n",
        "                'decoder': decoder.state_dict(),\n",
        "                'classifier': classifier.state_dict(),\n",
        "                'optimizer': optimizer.state_dict(),\n",
        "                'train_loss': mean_train_loss,\n",
        "                'train_acc': train_acc,\n",
        "                'valid_loss': valid_loss,\n",
        "                'valid_acc': valid_acc,\n",
        "                'epoch': epoch,\n",
        "                'timestamp': datetime.now().strftime(\"%y/%m/%d %H:%M:%S\"),\n",
        "            }\n",
        "            if train_params.c0 == 0:\n",
        "                model_dict['composite_loss'] = composite_loss.state_dict()\n",
        "\n",
        "            weights_path = '{}.{:d}'.format(train_params.weights_path, epoch)\n",
        "            print('Saving weights to {}'.format(weights_path))\n",
        "            torch.save(model_dict, weights_path)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  0%|          | 16384/26421880 [00:00<03:04, 143504.31it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Training on CPU\n",
            "Using Fashion MNIST dataset\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "26427392it [00:00, 76014930.78it/s]                            \n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "32768it [00:00, 492106.52it/s]\n",
            "  2%|▏         | 98304/4422102 [00:00<00:04, 947745.92it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz\n",
            "Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "4423680it [00:00, 23433810.85it/s]                         \n",
            "8192it [00:00, 148143.19it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw\n",
            "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz\n",
            "Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw\n",
            "Processing...\n",
            "Done!\n",
            "Training samples: 48000\n",
            "Validation samples: 12000\n",
            "Labeled training samples: 9600\n",
            "Loading weights from: /content/drive/My Drive/Colab Notebooks/checkpoints/fmnist/model.pt.373 with:\n",
            "\tepoch: 373\n",
            "\ttrain_loss: 0.580299197991689\n",
            "\ttrain_acc: 88.42083333333333\n",
            "\tvalid_loss: 1.6068398818969727\n",
            "\tvalid_acc: 89.45833333333333\n",
            "\ttimestamp: 19/11/22 19:18:54\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J8HMHUt5tfKQ",
        "colab_type": "code",
        "outputId": "8faa34e3-77b5-4565-9b86-34c7bd079ad0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        }
      },
      "source": [
        "print('Testing...')\n",
        "start_t = time.time()\n",
        "test_loss, test_acc, test_psnr = eval(\n",
        "    modules, test_dataloader, (criterion_rec, criterion_cls), device)\n",
        "end_t = time.time()\n",
        "test_time = end_t - start_t\n",
        "\n",
        "print('Test-Loss: %.6f | Test-Acc: %.3f%% | Test-PSNR: %.3f%% | Test-Time: %.3f sec' % (\n",
        "    test_loss, test_acc, test_psnr, test_time))"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Testing...\n",
            "Test-Loss: 1.636271 | Test-Acc: 86.680% | Test-PSNR: 139.279% | Test-Time: 6.049 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B3GG6KguzvQD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# fashion set:\n",
        "# 20 epochs 78.3%, stop and set c0 to 3\n",
        "\n",
        "\n",
        "#126 epoch 84%, lower lr 0.00005, change from 0.3 to 0.1\n",
        "#170 epoch 84.1% test acc, at epoch 170 set lr=0.00008\n",
        "#187 set c0=1, acc drop to 76%"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}